import sys
import os
import json
import requests
import pyaudio
import wave
import time
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout,
                             QHBoxLayout, QTextEdit, QPushButton, QLabel,
                             QComboBox, QProgressBar, QMessageBox)
from PyQt5.QtCore import Qt, QThread, pyqtSignal, QTimer
from PyQt5.QtGui import QIcon


class VoiceAIAssistant(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("语音AI助手")
        self.setGeometry(100, 100, 800, 600)

        # 初始化DeepSeek API设置
        self.api_url = "https://api.deepseek.com/v1/chat/completions"
        self.api_key = os.environ.get("DEEPSEEK_API_KEY")
        self.model = "deepseek-chat"

        # 音频参数
        self.CHUNK = 1024
        self.FORMAT = pyaudio.paInt16
        self.CHANNELS = 2
        self.RATE = 44100
        self.pyaudio = pyaudio.PyAudio()

        # 初始化UI
        self.init_ui()

        # 状态变量
        self.is_listening = False
        self.is_processing = False
        self.recording_start_time = 0
        self.audio_frames = []
        self.audio_stream = None

    def init_ui(self):
        # 创建主窗口部件
        main_widget = QWidget()
        self.setCentralWidget(main_widget)

        # 创建主布局
        main_layout = QVBoxLayout(main_widget)
        main_layout.setContentsMargins(15, 15, 15, 15)

        # 标题
        title_label = QLabel("语音AI助手")
        title_label.setStyleSheet("font-size: 24px; font-weight: bold; color: #2c3e50;")
        title_label.setAlignment(Qt.AlignCenter)
        main_layout.addWidget(title_label)

        # 输入输出区域
        io_layout = QVBoxLayout()

        # 输入框
        input_label = QLabel("用户输入:")
        input_label.setStyleSheet("font-size: 16px; font-weight: bold;")
        io_layout.addWidget(input_label)

        self.input_text = QTextEdit()
        self.input_text.setPlaceholderText("点击下方麦克风按钮说话，或直接在此输入文本...")
        self.input_text.setStyleSheet("font-size: 14px; min-height: 100px;")
        io_layout.addWidget(self.input_text)

        # 输出框
        output_label = QLabel("AI回复:")
        output_label.setStyleSheet("font-size: 16px; font-weight: bold; margin-top: 15px;")
        io_layout.addWidget(output_label)

        self.output_text = QTextEdit()
        self.output_text.setReadOnly(True)
        self.output_text.setPlaceholderText("AI回复将显示在这里...")
        self.output_text.setStyleSheet("font-size: 14px; min-height: 200px;")
        io_layout.addWidget(self.output_text)

        main_layout.addLayout(io_layout)

        # 控制区域
        control_layout = QHBoxLayout()

        # 麦克风按钮
        self.mic_button = QPushButton()
        self.mic_button.setIcon(QIcon.fromTheme("microphone"))
        self.mic_button.setStyleSheet("""
            QPushButton {
                font-size: 16px;
                padding: 10px;
                min-width: 100px;
                border-radius: 5px;
                background-color: #e74c3c;
                color: white;
            }
            QPushButton:hover {
                background-color: #c0392b;
            }
        """)
        self.mic_button.setText("开始录音")
        self.mic_button.clicked.connect(self.toggle_listening)
        control_layout.addWidget(self.mic_button)

        # 录音计时标签
        self.recording_time_label = QLabel("00:00")
        self.recording_time_label.setStyleSheet("font-size: 16px; min-width: 60px;")
        self.recording_time_label.setAlignment(Qt.AlignCenter)
        self.recording_time_label.hide()
        control_layout.addWidget(self.recording_time_label)

        # 提交按钮
        self.submit_button = QPushButton("提交问题")
        self.submit_button.setStyleSheet("""
            QPushButton {
                font-size: 16px;
                padding: 10px;
                min-width: 100px;
                border-radius: 5px;
                background-color: #3498db;
                color: white;
            }
            QPushButton:hover {
                background-color: #2980b9;
            }
        """)
        self.submit_button.clicked.connect(self.submit_question)
        control_layout.addWidget(self.submit_button)

        # 模型选择
        model_label = QLabel("模型:")
        model_label.setStyleSheet("font-size: 14px;")
        control_layout.addWidget(model_label)

        self.model_combo = QComboBox()
        self.model_combo.addItems(["deepseek-chat", "deepseek-coder"])
        self.model_combo.setCurrentText(self.model)
        self.model_combo.setStyleSheet("font-size: 14px; padding: 5px;")
        control_layout.addWidget(self.model_combo)

        main_layout.addLayout(control_layout)

        # 进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setTextVisible(False)
        self.progress_bar.setRange(0, 0)  # 不确定模式
        self.progress_bar.setFixedHeight(5)
        self.progress_bar.setVisible(False)
        main_layout.addWidget(self.progress_bar)

        # 状态栏
        self.status_bar = QLabel("准备就绪")
        self.status_bar.setStyleSheet("font-size: 12px; color: #7f8c8d;")
        main_layout.addWidget(self.status_bar)

        # 录音计时器
        self.recording_timer = QTimer()
        self.recording_timer.timeout.connect(self.update_recording_time)

    def toggle_listening(self):
        """切换录音状态"""
        if self.is_listening:
            self.stop_listening()
        else:
            self.start_listening()

    def start_listening(self):
        """开始录音"""
        if self.is_processing:
            QMessageBox.warning(self, "警告", "正在处理上一个请求，请稍候...")
            return

        # 列出音频设备并选择默认输入设备
        device_index = self.get_default_input_device()
        if device_index is None:
            QMessageBox.warning(self, "警告", "找不到可用的音频输入设备")
            return

        try:
            # 打开音频流
            self.audio_stream = self.pyaudio.open(
                input_device_index=device_index,
                format=self.FORMAT,
                channels=self.CHANNELS,
                rate=self.RATE,
                input=True,
                frames_per_buffer=self.CHUNK
            )

            self.is_listening = True
            self.audio_frames = []
            self.recording_start_time = time.time()

            self.mic_button.setText("停止录音")
            self.mic_button.setStyleSheet("""
                QPushButton {
                    font-size: 16px;
                    padding: 10px;
                    min-width: 100px;
                    border-radius: 5px;
                    background-color: #2ecc71;
                    color: white;
                }
                QPushButton:hover {
                    background-color: #27ae60;
                }
            """)

            self.recording_time_label.setText("00:00")
            self.recording_time_label.show()
            self.status_bar.setText("正在录音...请说话")

            # 开始录音计时器
            self.recording_timer.start(200)  # 每200ms更新一次时间

            # 在后台线程中录音
            self.recording_thread = RecordingThread(
                audio_stream=self.audio_stream,
                chunk=self.CHUNK,
                frames=self.audio_frames
            )
            self.recording_thread.finished.connect(self.handle_recording_result)
            self.recording_thread.error.connect(self.handle_recording_error)
            self.recording_thread.start()

        except Exception as e:
            QMessageBox.warning(self, "错误", f"无法开始录音: {str(e)}")
            self.stop_listening()

    def stop_listening(self):
        """停止录音"""
        if not self.is_listening:
            return

        self.is_listening = False

        # 停止录音线程
        if hasattr(self, 'recording_thread') and self.recording_thread.isRunning():
            self.recording_thread.stop()
            self.recording_thread.wait()

        # 停止音频流
        if self.audio_stream:
            self.audio_stream.stop_stream()
            self.audio_stream.close()
            self.audio_stream = None

        # 停止计时器
        self.recording_timer.stop()

        self.mic_button.setText("开始录音")
        self.mic_button.setStyleSheet("""
            QPushButton {
                font-size: 16px;
                padding: 10px;
                min-width: 100px;
                border-radius: 5px;
                background-color: #e74c3c;
                color: white;
            }
            QPushButton:hover {
                background-color: #c0392b;
            }
        """)
        self.recording_time_label.hide()
        self.status_bar.setText("录音已停止")

    def update_recording_time(self):
        """更新录音时间显示"""
        elapsed = int(time.time() - self.recording_start_time)
        minutes = elapsed // 60
        seconds = elapsed % 60
        self.recording_time_label.setText(f"{minutes:02d}:{seconds:02d}")

    def get_default_input_device(self):
        """获取默认输入设备索引"""
        for i in range(self.pyaudio.get_device_count()):
            dev = self.pyaudio.get_device_info_by_index(i)
            if dev['maxInputChannels'] > 0:
                return i
        return None

    def handle_recording_result(self):
        """处理录音结果"""
        if not self.audio_frames:
            self.status_bar.setText("没有检测到语音")
            return

        # 保存录音文件
        filename = f"recording_{int(time.time())}.wav"
        try:
            with wave.open(filename, 'wb') as wf:
                wf.setnchannels(self.CHANNELS)
                wf.setsampwidth(self.pyaudio.get_sample_size(self.FORMAT))
                wf.setframerate(self.RATE)
                wf.writeframes(b''.join(self.audio_frames))

            self.status_bar.setText(f"录音已保存: {filename}")

            # 在这里可以添加语音识别代码
            # 暂时直接将录音信息显示在输入框
            duration = time.time() - self.recording_start_time
            self.input_text.setPlainText(f"[录音时长: {duration:.1f}秒] 请点击提交按钮进行语音识别")

        except Exception as e:
            QMessageBox.warning(self, "错误", f"保存录音文件失败: {str(e)}")
            self.status_bar.setText("录音保存失败")

    def handle_recording_error(self, error):
        """处理录音错误"""
        self.stop_listening()
        QMessageBox.warning(self, "错误", f"录音过程中发生错误: {error}")
        self.status_bar.setText("录音错误")

    def submit_question(self):
        """提交问题给AI"""
        question = self.input_text.toPlainText().strip()
        if not question:
            QMessageBox.warning(self, "警告", "请输入问题或使用语音输入")
            return

        if self.is_processing:
            QMessageBox.warning(self, "警告", "正在处理上一个请求，请稍候...")
            return

        self.is_processing = True
        self.progress_bar.setVisible(True)
        self.status_bar.setText("正在处理请求...")

        # 更新模型选择
        self.model = self.model_combo.currentText()

        # 在后台线程中处理请求
        self.ai_thread = AIThread(
            api_url=self.api_url,
            api_key=self.api_key,
            model=self.model,
            question=question
        )
        self.ai_thread.finished.connect(self.handle_ai_response)
        self.ai_thread.error.connect(self.handle_ai_error)
        self.ai_thread.start()

    def handle_ai_response(self, response):
        """处理AI回复"""
        self.output_text.setHtml(self.format_response(response))
        self.is_processing = False
        self.progress_bar.setVisible(False)
        self.status_bar.setText("请求完成")

    def handle_ai_error(self, error):
        """处理AI错误"""
        self.is_processing = False
        self.progress_bar.setVisible(False)
        QMessageBox.critical(self, "错误", f"请求失败: {error}")
        self.status_bar.setText("请求失败")

    def format_response(self, text):
        """格式化AI回复为HTML"""
        # 简单的Markdown转HTML
        text = text.replace("\n", "<br>")
        text = text.replace("**", "<strong>", 1)
        text = text.replace("**", "</strong>", 1)
        text = text.replace("*", "<em>", 1)
        text = text.replace("*", "</em>", 1)

        return f"""
        <div style="font-family: Arial, sans-serif; line-height: 1.6;">
            {text}
        </div>
        """

    def closeEvent(self, event):
        """关闭窗口时确保所有线程停止"""
        self.stop_listening()

        if hasattr(self, 'ai_thread') and self.ai_thread.isRunning():
            self.ai_thread.wait()

        self.pyaudio.terminate()
        event.accept()


class RecordingThread(QThread):
    """录音线程"""
    finished = pyqtSignal()
    error = pyqtSignal(str)

    def __init__(self, audio_stream, chunk, frames):
        super().__init__()
        self.audio_stream = audio_stream
        self.chunk = chunk
        self.frames = frames
        self._is_running = True

    def run(self):
        try:
            while self._is_running:
                data = self.audio_stream.read(self.chunk, exception_on_overflow=False)
                self.frames.append(data)

            self.finished.emit()
        except Exception as e:
            self.error.emit(str(e))

    def stop(self):
        self._is_running = False


class AIThread(QThread):
    """AI请求线程"""
    finished = pyqtSignal(str)
    error = pyqtSignal(str)

    def __init__(self, api_url, api_key, model, question):
        super().__init__()
        self.api_url = api_url
        self.api_key = api_key
        self.model = model
        self.question = question

    def run(self):
        try:
            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {self.api_key}"
            }

            payload = {
                "model": self.model,
                "messages": [{"role": "user", "content": self.question}],
                "temperature": 0.7,
                "max_tokens": 2000
            }

            response = requests.post(
                self.api_url,
            else:
                self.error.emit(f"API错误: {response.status_code} - {response.text}")
        except Exception as e:
            self.error.emit(str(e))


def main():
    app = QApplication(sys.argv)
    window = VoiceAIAssistant()
    window.show()
    sys.exit(app.exec_())


if __name__ == "__main__":
    main()